#!/usr/bin/env python
import os
import argparse
import logging
import math
import matplotlib.pyplot as plot
import pandas as pd
from rosbags.rosbag2 import Reader
from rosbags.serde import deserialize_cdr
from pathlib import Path
from functools import reduce
from tqdm import tqdm
import numpy as np
import pyproj
import datetime
from scipy.spatial.transform import Rotation
from numpy.linalg import inv

# 全局配置
ISOTIMEFORMAT = '%Y-%m-%d_%H:%M:%S'
theTime_ = datetime.datetime.now().strftime(ISOTIMEFORMAT)
glob_save_st_eval_pic_flag_ = False
st_eval_flag_ = 'eval_'  # 评测结果前缀
eva_dis_list_ = []
eva_yaw_list_ = []
all_result_path = ""  # 评测结果根目录

class InsType:
    """INS消息数据索引定义"""
    def __init__(self):
        self.utc_time_us_index = 0
        self.sys_time_us_index = 1
        self.status_index = 2
        self.lat_index = 3
        self.lon_index = 4
        self.height_index = 5
        self.latitude_standard_index = 6
        self.longitude_standard_index = 7
        self.height_standard_index = 8
        self.heading_index = 9


class VehicleType:
    """车辆速度消息索引定义"""
    def __init__(self):
        self.stamp_sec = 0
        self.stamp_nanosec = 1
        self.vehicle_speed_ts = 2
        self.vehicle_speed = 3


class OdomType:
    """DR里程计消息索引定义"""
    def __init__(self):
        self.stamp_sec_index = 0
        self.stamp_nanosec_index = 1
        self.position_x_index = 2
        self.position_y_index = 3
        self.yaw_index = 4


def Euler2RotDegree(euler):
    """欧拉角(度)转旋转矩阵"""
    r = Rotation.from_euler('zyx', euler, degrees=True)
    return r.as_matrix()


def rotation_matrix_to_euler_angles(R):
    """旋转矩阵转欧拉角(弧度)"""
    assert R.shape == (3, 3), "旋转矩阵必须是3x3"
    assert np.allclose(np.dot(R.T, R), np.eye(3)), "旋转矩阵必须正交"

    sy = np.sqrt(R[0, 0] **2 + R[1, 0]** 2)
    singular = sy < 1e-6

    if not singular:
        x = np.arctan2(R[2, 1], R[2, 2])
        y = np.arctan2(-R[2, 0], sy)
        z = np.arctan2(R[1, 0], R[0, 0])
    else:
        x = np.arctan2(-R[1, 2], R[1, 1])
        y = np.arctan2(-R[2, 0], sy)
        z = 0
    return np.array([z, y, x])  # z-y-x顺序


class EvaluationProcessor(object):
    """DR与INS精度对比处理器"""
    TOPICS = {
        '/SA/INSPVA': {
            'columns': ['utc_time_us', 'sys_time_us', 'ins_status', 'latitude', 'longitude', 'heigth',
                        'latitude_standard_deviation', 'longitude_standard_deviation', 'height_standard_deviation',
                        'heading', 'ins_speed', 'east_velocity', 'north_velocity', 'up_velocity',
                        'east_velocity_standard_deviation', 'north_velocity_standard_deviation',
                        'up_velocity_standard_deviation', 'status','accel_x']
        },
        '/VA/VehicleMotionIpd': {
            'columns': ['header.stamp.sec', 'header.stamp.nanosec', 'vehicle_speed_ts', 'vehicle_speed']
        },
        '/PK/DR/Result': {
            'columns': ['header.stamp.sec', 'header.stamp.nanosec', 'position_x', 'position_y', 'headrate']
        }
    }

    def __init__(self, msg_path: str, bag_path: str, use_source_mode: bool = False):
        super().__init__()
        # 存储INS和DR的评估坐标与姿态
        self.eva_ins_x = []
        self.eva_ins_y = []
        self.eva_ins_yaw = []
        self.eva_odom_x = []
        self.eva_odom_y = []
        self.eva_odom_yaw = []
        
        self.rad2degree = 180 / np.pi
        self.odom_mileage = 0.0  # DR里程
        self.ins_mileage = 0.0   # INS里程

        # 初始化ROS2 bag读取器
        from msg_register import Register
        Register().register_msg_types(msg_path, use_source_mode=use_source_mode)  # 注册消息类型
        self.bag_reader = Reader(bag_path)
        self.bag_reader.open()
        
        # 统计需要处理的消息数量
        self.total_msgs = sum(
            1 for connection, _, _ in self.bag_reader.messages() if connection.topic in self.TOPICS.keys()
        )

    def get_last_attribute(self, obj, attribute_path):
        """获取嵌套属性的值"""
        attributes = attribute_path.split('.')
        return reduce(getattr, attributes, obj)

    def _timestamp_synchron(self, odom_data, ins_data):
        """时间戳同步：对齐ODOM和INS的起始时间"""
        if not odom_data or not ins_data:
            print("错误：ODOM或INS数据为空，无法同步时间戳")
            return -1, odom_data, ins_data

        # ODOM起始时间（秒）
        odom_first_ts = odom_data[0][OdomType().stamp_sec_index] + \
                       odom_data[0][OdomType().stamp_nanosec_index] * 1e-9
        
        # 过滤无效ODOM数据（时间戳接近0的）
        while odom_data and (odom_data[0][OdomType().stamp_sec_index] + 
                            odom_data[0][OdomType().stamp_nanosec_index] * 1e-9) < 0.1:
            odom_data.pop(0)

        # 寻找INS中与ODOM起始时间最接近的点
        ins_ts_list = [data[InsType().sys_time_us_index] * 1e-6 for data in ins_data]
        min_diff = float('inf')
        best_idx = 0
        for i, ts in enumerate(ins_ts_list):
            diff = abs(ts - odom_first_ts)
            if diff < min_diff and diff < 0.2:  # 允许200ms内的误差
                min_diff = diff
                best_idx = i

        # 截断INS数据至同步点
        if best_idx > 0:
            ins_data = ins_data[best_idx:]
        
        print(f"时间同步完成：ODOM起始时间={odom_first_ts:.3f}s, INS起始时间={ins_ts_list[best_idx]:.3f}s")
        return 1, odom_data, ins_data

    def _ins_lla2enu(self, ins_data):
        """将INS的LLA坐标转换为ENU坐标系（相对起始点）"""
        if not ins_data:
            print("错误：INS数据为空，无法转换坐标")
            return -1

        # 以第一个有效点为原点
        start_lon = ins_data[0][InsType().lon_index]
        start_lat = ins_data[0][InsType().lat_index]
        start_heading = ins_data[0][InsType().heading_index]
        start_heading = start_heading - 360 if start_heading > 360 else start_heading

        # 初始化ENU投影
        proj = pyproj.Proj(proj='tmerc', lon_0=start_lon, lat_0=start_lat, preserve_units=False)
        start_enu_x, start_enu_y = proj(start_lon, start_lat)

        # 转换所有INS点到ENU
        for data in ins_data:
            lon = data[InsType().lon_index]
            lat = data[InsType().lat_index]
            heading = data[InsType().heading_index]
            
            # 计算ENU坐标（相对原点）
            enu_x, enu_y = proj(lon, lat)
            # 转换航向角（ENU坐标系下）
            enu_yaw = -heading + 90.0  # 转换为ENU坐标系下的偏航角
            
            self.eva_ins_x.append(enu_x)
            self.eva_ins_y.append(enu_y)
            self.eva_ins_yaw.append(enu_yaw)

        print(f"INS坐标转换完成：共{len(ins_data)}个点，起始ENU坐标=({start_enu_x:.2f}, {start_enu_y:.2f})")
        return 1

    def _odom_to_enu(self, odom_data):
        """将DR里程计坐标转换到INS的ENU坐标系"""
        if not odom_data or not self.eva_ins_x:
            print("错误：ODOM数据为空或INS坐标未初始化")
            return -1

        # 以第一个ODOM点为起始点
        start_odom_x = odom_data[0][OdomType().position_x_index]
        start_odom_y = odom_data[0][OdomType().position_y_index]
        start_odom_yaw = odom_data[0][OdomType().yaw_index]

        # 计算ODOM到ENU的转换矩阵
        odom_euler = [start_odom_yaw, 0.0, 0.0]
        rot_odom = Euler2RotDegree(odom_euler)
        # ODOM到车辆坐标系的逆变换
        t_odom_to_vehicle = inv(np.array([
            [rot_odom[0,0], rot_odom[0,1], rot_odom[0,2], start_odom_x],
            [rot_odom[1,0], rot_odom[1,1], rot_odom[1,2], start_odom_y],
            [rot_odom[2,0], rot_odom[2,1], rot_odom[2,2], 0.0],
            [0.0, 0.0, 0.0, 1.0]
        ]))

        # 车辆到ENU的变换（使用INS起始点）
        ins_euler = [self.eva_ins_yaw[0], 0.0, 0.0]
        rot_ins = Euler2RotDegree(ins_euler)
        t_vehicle_to_enu = np.array([
            [rot_ins[0,0], rot_ins[0,1], rot_ins[0,2], self.eva_ins_x[0]],
            [rot_ins[1,0], rot_ins[1,1], rot_ins[1,2], self.eva_ins_y[0]],
            [rot_ins[2,0], rot_ins[2,1], rot_ins[2,2], 0.0],
            [0.0, 0.0, 0.0, 1.0]
        ])

        # 组合变换矩阵：ODOM → 车辆 → ENU
        t_odom_to_enu = np.dot(t_vehicle_to_enu, t_odom_to_vehicle)

        # 转换所有ODOM点到ENU
        for data in odom_data:
            x = data[OdomType().position_x_index]
            y = data[OdomType().position_y_index]
            yaw = data[OdomType().yaw_index]

            # 转换位置
            odom_pose = np.array([[x], [y], [0.0], [1.0]])
            enu_pose = np.dot(t_odom_to_enu, odom_pose)

            # 转换姿态
            rot_current = Euler2RotDegree([yaw, 0.0, 0.0])
            t_current_vehicle_to_odom = np.array([
                [rot_current[0,0], rot_current[0,1], rot_current[0,2], x],
                [rot_current[1,0], rot_current[1,1], rot_current[1,2], y],
                [rot_current[2,0], rot_current[2,1], rot_current[2,2], 0.0],
                [0.0, 0.0, 0.0, 1.0]
            ])
            t_enu = np.dot(t_odom_to_enu, t_current_vehicle_to_odom)
            ypr_deg = np.degrees(rotation_matrix_to_euler_angles(t_enu[:3,:3]))

            self.eva_odom_x.append(t_enu[0,3])
            self.eva_odom_y.append(t_enu[1,3])
            self.eva_odom_yaw.append(ypr_deg[0])

        print(f"ODOM坐标转换完成：共{len(odom_data)}个degrees点")
        return 1

    def _calculate_errors(self, odom_data):
        """计算DR与INS的误差"""
        # if len(self.eva_odom_x) != len(self.eva_ins_x):
        #     print(f"警告：ODOM与INS数据点数量不匹配（{len(self.eva_odom_x)} vs {len(self.eva_ins_x)}）")

        # 计算里程
        for i in range(1, len(self.eva_odom_x)):
            dx = self.eva_odom_x[i] - self.eva_odom_x[i-1]
            dy = self.eva_odom_y[i] - self.eva_odom_y[i-1]
            self.odom_mileage += math.hypot(dx, dy)
        
        for i in range(1, len(self.eva_ins_x)):
            dx = self.eva_ins_x[i] - self.eva_ins_x[i-1]
            dy = self.eva_ins_y[i] - self.eva_ins_y[i-1]
            self.ins_mileage += math.hypot(dx, dy)

        # 计算终点误差
        # 方式1、ENU坐标系下的欧式误差
        # dx = self.eva_ins_x[-1] - self.eva_odom_x[-1]
        # dy = self.eva_ins_y[-1] - self.eva_odom_y[-1]
        # dis_error = math.hypot(dx, dy)

        # 方式2、车辆坐标系下的纵向、横向误差
        # print(f"终点yaw角: {self.eva_ins_yaw[-1]}")
        rot_ins = Euler2RotDegree([self.eva_ins_yaw[-1], 0.0, 0.0])
        t_enu2vehicle = inv(np.array([
            [rot_ins[0,0], rot_ins[0,1], 0, self.eva_ins_x[-1]],
            [rot_ins[1,0], rot_ins[1,1], 0, self.eva_ins_y[-1]],
            [0, 0, 1, 0],
            [0, 0, 0, 1]
        ]))
        odom_pose = np.array([[self.eva_odom_x[-1]], [self.eva_odom_y[-1]], [0], [1]])
        vehicle_frame = np.dot(t_enu2vehicle, odom_pose)
        dx = vehicle_frame[0][0]
        dy = vehicle_frame[1][0]
        dis_error = math.hypot(dx, dy)

        # 3、航向误差
        yaw_bias = self.eva_ins_yaw[0] - self.eva_odom_yaw[0] # 初始航向误差
        eva_odom_yaw_corrected = [y + yaw_bias for y in self.eva_odom_yaw] # 航向误差修正
        yaw_error = abs(self.eva_ins_yaw[-1] - eva_odom_yaw_corrected[-1]) # 终点航向误差
        print(f"终点误差计算完成：距离误差={dis_error:.4f}m, 航向误差={yaw_error:.4f}°")
        print(f"x, y误差: {dx}m, {dy}m")
        print(f"里程计算完成: INS里程={self.ins_mileage:.2f}m, ODOM里程={self.odom_mileage:.2f}m")
        
        if yaw_error > 180:
            yaw_error = 360 - yaw_error

        # 计算误差率
        mileage = self.odom_mileage if self.odom_mileage > 0.01 else 1.0
        rel_error = (dis_error / mileage) * 100 if mileage > 0 else 0
        rel_x_error = (abs(dx) / mileage) * 100 if mileage > 0 else 0
        rel_y_error = (abs(dy) / mileage) * 100 if mileage > 0 else 0

        print(f"相对误差计算完成：距离误差={rel_error:.4f}%")
        return [dis_error, rel_error, mileage, dx, dy, yaw_error, rel_x_error, rel_y_error]

    def process(self, output_path, save_pic=True):
        """主处理流程"""
        # 提取消息数据
        ins_data = []
        odom_data = []
        vehicle_data = []

        for connection, _, rawdata in tqdm(
            self.bag_reader.messages(), desc="解析Bag数据", total=self.total_msgs
        ):
            if connection.topic not in self.TOPICS:
                continue

            try:
                msg = deserialize_cdr(rawdata, connection.msgtype)
            except AssertionError:
                print(f"跳过损坏的消息（主题：{connection.topic}）")
                continue

            # 提取消息字段
            row = [
                self.get_last_attribute(msg, col) 
                for col in self.TOPICS[connection.topic]['columns']
            ]

            # 分类存储数据
            if connection.topic == '/SA/INSPVA':
                ins_data.append([
                    row[InsType().utc_time_us_index],
                    row[InsType().sys_time_us_index],
                    row[InsType().status_index],
                    row[InsType().lat_index],
                    row[InsType().lon_index],
                    row[InsType().height_index],
                    row[InsType().longitude_standard_index],
                    row[InsType().latitude_standard_index],
                    row[InsType().height_standard_index],
                    row[InsType().heading_index]
                ])
            elif connection.topic == '/VA/VehicleMotionIpd':
                vehicle_data.append([
                    row[VehicleType().stamp_sec],
                    row[VehicleType().stamp_nanosec],
                    row[VehicleType().vehicle_speed_ts],
                    row[VehicleType().vehicle_speed]
                ])
            elif connection.topic == '/PK/DR/Result':
                odom_data.append([
                    row[OdomType().stamp_sec_index],
                    row[OdomType().stamp_nanosec_index],
                    row[OdomType().position_x_index],
                    row[OdomType().position_y_index],
                    row[OdomType().yaw_index] * self.rad2degree  # 弧度转角度
                ])

        # 数据处理流程
        ret, odom_data, ins_data = self._timestamp_synchron(odom_data, ins_data)
        if ret < 0:
            return -1

        if self._ins_lla2enu(ins_data) < 0:
            return -1

        if self._odom_to_enu(odom_data) < 0:
            return -1

        # 计算误差
        eva_result = self._calculate_errors(odom_data)
        if not eva_result:
            return -1

        # 保存结果
        self._save_results(output_path, odom_data, ins_data, eva_result)
        
        # 生成并保存轨迹图
        if save_pic:
            self._plot_trajectory(output_path, eva_result)

        return 1

    def _save_results(self, output_path, odom_data, ins_data, eva_result):
        """保存评测结果到CSV文件"""
        Path(output_path).mkdir(parents=True, exist_ok=True)
        
        # 保存ODOM评测数据
        odom_df = pd.DataFrame({
            'stamp_sec': [d[OdomType().stamp_sec_index] for d in odom_data],
            'stamp_nanosec': [d[OdomType().stamp_nanosec_index] for d in odom_data],
            'enu_x': self.eva_odom_x,
            'enu_y': self.eva_odom_y,
            'yaw_deg': self.eva_odom_yaw
        })
        odom_df.to_csv(Path(output_path)/'odom_evaluation.csv', index=False, float_format='%.4f')

        # 保存INS评测数据
        ins_df = pd.DataFrame({
            'utc_time_us': [d[InsType().utc_time_us_index] for d in ins_data],
            'sys_time_us': [d[InsType().sys_time_us_index] for d in ins_data],
            'enu_x': self.eva_ins_x,
            'enu_y': self.eva_ins_y,
            'yaw_deg': self.eva_ins_yaw,
            'status': [d[InsType().status_index] for d in ins_data]
        })
        ins_df.to_csv(Path(output_path)/'ins_evaluation.csv', index=False, float_format='%.4f')

        # 保存误差结果
        error_df = pd.DataFrame({
            '误差距离(m)': [eva_result[0]],
            '相对误差(%)': [eva_result[1]],
            '行驶里程(m)': [eva_result[2]],
            'X方向误差(m)': [eva_result[3]],
            'Y方向误差(m)': [eva_result[4]],
            '航向误差(度)': [eva_result[5]],
            'X相对误差(%)': [eva_result[6]],
            'Y相对误差(%)': [eva_result[7]]
        })
        error_df.to_csv(Path(output_path)/'error_evaluation.csv', index=False, float_format='%.4f')

        # 全局结果汇总
        global eva_dis_list_, eva_yaw_list_, all_result_path
        eva_dis_list_.append(eva_result[0])
        eva_yaw_list_.append(eva_result[5])

    def _plot_trajectory(self, output_path, eva_result):
        """绘制DR与INS轨迹对比图"""
        plot.figure(figsize=(10, 8))
        
        # 绘制轨迹
        plot.plot(self.eva_odom_x, self.eva_odom_y, 
                label='DR Trajectory', c='r', linestyle='--', linewidth=0.8)
        plot.plot(self.eva_ins_x, self.eva_ins_y, 
                label='INS Trajectory', c='g', linestyle='-', linewidth=0.8)
        
        # 图形设置
        plot.xlabel('X Coordinate (m)')
        plot.ylabel('Y Coordinate (m)')
        plot.title('DR vs INS Trajectory Comparison')
        plot.legend()
        plot.grid(True)
        plot.axis('equal')
        
        # 添加误差信息
        info_text = (f"Error Distance: {eva_result[0]:.4f}m\n"
                    f"Relative Error: {eva_result[1]:.2f}%\n"
                    f"X Error: {eva_result[3]:.4f}m\n"  # X方向误差
                    f"Y Error: {eva_result[4]:.4f}m\n"  # Y方向误差
                    f"Yaw Error: {eva_result[5]:.4f}°\n"
                    f"Travel Distance: {eva_result[2]:.2f}m")
        plot.text(0.05, 0.95, info_text, transform=plot.gca().transAxes,
                verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        
        # 保存图片
        plot.savefig(Path(output_path)/'trajectory_comparison.png', dpi=300, bbox_inches='tight')
        plot.close()
        print("Trajectory comparison plot has been saved")

    def __del__(self):
        """关闭bag读取器"""
        try:
            self.bag_reader.close()
        except Exception as e:
            print(f"关闭bag时出错: {e}")


def main():
    parser = argparse.ArgumentParser(description='DR与INS精度对比评测工具')
    parser.add_argument('-m', '--msg-path', type=str, required=True, help='ROS消息定义文件路径')
    parser.add_argument('-b', '--bag-path', type=str, help='单个rosbag文件路径')
    parser.add_argument('-o', '--output-path', type=str, required=True, help='评测结果输出路径')
    parser.add_argument('-p', '--pic', type=bool, default=True, help='是否保存轨迹图')
    parser.add_argument('-d', '--dir', type=str, help='rosbag目录（批量处理）')
    parser.add_argument('-n', '--name', type=str, required=True, help='bag名称关键词')
    parser.add_argument('-st', '--model', type=str, default='eval', help='评测模式 (eval/st)')
    parser.add_argument(
        '--use-source-mode',
        action='store_true',
        help='使用ROS2 msg源文件模式（默认使用install目录模式）'
    )

    args = parser.parse_args()
    global all_result_path
    all_result_path = args.output_path

    # 批量处理模式
    if args.dir:
        Path(args.output_path).mkdir(parents=True, exist_ok=True)
        
        # 查找符合名称关键词的bag文件或目录
        # print(f"查找包含关键词 '{args.name}' 的bag文件或目录, 输出路径: {args.output_path},.dir: {args.dir}")
        bag_dirs = []
        for root, dirs, _ in os.walk(args.dir):
            for d in dirs:
                if args.name in d:            # 目录名里包含关键词
                    bag_dir = Path(root) / d
                    # 再确认目录里至少有一个 .db3
                    if any(f.suffix == '.db3' for f in bag_dir.iterdir()):
                        bag_dirs.append(bag_dir)
                        print(f"找到bag目录: {bag_dir}")
        
        if not bag_dirs:
            print(f"未找到包含关键词 '{args.name}' 的bag文件")
            return

        # 处理每个bag文件
        for bag_path in bag_dirs:
            print(f"\n===== 处理bag文件: {bag_path} =====")
            output_dir = Path(args.output_path)/bag_path.stem
            processor = EvaluationProcessor(args.msg_path, str(bag_path), use_source_mode=args.use_source_mode)
            processor.process(str(output_dir), args.pic)

        # 计算统计结果
        if eva_dis_list_ and eva_yaw_list_:
            eva_dis_list_.sort()
            eva_yaw_list_.sort()
            total = len(eva_dis_list_)
            
            # 计算统计指标
            avg_dis = sum(eva_dis_list_) / total
            avg_yaw = sum(eva_yaw_list_) / total
            cep68 = eva_dis_list_[int(total*0.68)] if total > 0 else 0
            cep95 = eva_dis_list_[int(total*0.95)] if total > 0 else 0

            # 保存统计结果
            stats_df = pd.DataFrame({
                '统计指标': ['平均距离误差(m)', '平均航向误差(度)', '68%置信距离误差(m)', '95%置信距离误差(m)', '样本数量'],
                '数值': [avg_dis, avg_yaw, cep68, cep95, total]
            })
            stats_df.to_csv(Path(args.output_path)/'statistics.csv', index=False, float_format='%.4f')
            print(f"\n统计结果已保存: {Path(args.output_path)/'statistics.csv'}")

    # 单个文件处理模式
    elif args.bag_path:
        output_dir = Path(args.output_path)
        processor = EvaluationProcessor(args.msg_path, args.bag_path, use_source_mode=args.use_source_mode)
        processor.process(str(output_dir), args.pic)

    else:
        print("错误: 必须指定--bag-path（单个文件）或--dir（批量处理）")


if __name__ == '__main__':
    main()
    